#include <vtkActor.h>
#include <vtkCallbackCommand.h>
#include <vtkCamera.h>
#include <vtkCellData.h>
#include <vtkCellPicker.h>
#include <vtkHardwarePicker.h>
#include <vtkCommand.h>
#include <vtkDataSetMapper.h>
#include <vtkExtractGeometry.h>
#include <vtkHexahedron.h>
#include <vtkIdTypeArray.h>
#include <vtkInteractorStyleTrackballCamera.h>
#include <vtkLookupTable.h>
#include <vtkNamedColors.h>
#include <vtkNew.h>
#include <vtkPlane.h>
#include <vtkPlaneWidget.h>
#include <vtkPoints.h>
#include <vtkProperty.h>
#include <vtkRenderWindow.h>
#include <vtkRenderWindowInteractor.h>
#include <vtkRenderer.h>
#include <vtkTetra.h>
#include <vtkUnsignedCharArray.h>
#include <vtkUnstructuredGrid.h>
#include <vtkSelectionNode.h>
#include <vtkSelection.h>
#include <vtkExtractSelection.h>
#include <vtkPolyDataMapper.h>
#include <vtkGeometryFilter.h>
#include <vtkSMPTools.h>

#include <iostream>

// 全局变量
vtkSmartPointer<vtkUnstructuredGrid> originalGrid;
vtkSmartPointer<vtkUnstructuredGrid> selectedGrid;
vtkSmartPointer<vtkExtractGeometry> extractGeometry;
vtkSmartPointer<vtkExtractGeometry> extractGeometrySelected;
vtkSmartPointer<vtkPolyDataMapper> mapper;
vtkSmartPointer<vtkDataSetMapper> selectedMapper;
vtkSmartPointer<vtkActor> actor;
vtkSmartPointer<vtkActor> selectedActor;
vtkSmartPointer<vtkRenderer> renderer;
vtkSmartPointer<vtkPlane> plane;
vtkSmartPointer<vtkPlaneWidget> planeWidget;

// 创建立方体网格（扩展版本，包含更多单元）
vtkSmartPointer<vtkUnstructuredGrid> CreateCubeGrid()
{
    vtkNew<vtkUnstructuredGrid> ugrid;
    vtkNew<vtkPoints> points;
    
    // 创建一个3x3x3的六面体网格
    int dimX = 4, dimY = 4, dimZ = 4;

    // 创建点
    for (int k = 0; k < dimZ; k++) {
        for (int j = 0; j < dimY; j++) {
            for (int i = 0; i < dimX; i++) {
                points->InsertNextPoint(i * 2.0, j * 2.0, k * 2.0);
            }
        }
    }

    ugrid->SetPoints(points);

    // 计算总的单元数量
    int numCells = (dimX - 1) * (dimY - 1) * (dimZ - 1);
    
    // 预分配内存以提高性能
    ugrid->Allocate(numCells);

    for (int k = 0; k < dimZ - 1; k++) {
        for (int j = 0; j < dimY - 1; j++) {
            for (int i = 0; i < dimX - 1; i++) {
                vtkNew<vtkHexahedron> hex;

                // 计算八个顶点的索引
                int idx000 = k * dimX * dimY + j * dimX + i;
                int idx100 = k * dimX * dimY + j * dimX + (i + 1);
                int idx110 = k * dimX * dimY + (j + 1) * dimX + (i + 1);
                int idx010 = k * dimX * dimY + (j + 1) * dimX + i;
                int idx001 = (k + 1) * dimX * dimY + j * dimX + i;
                int idx101 = (k + 1) * dimX * dimY + j * dimX + (i + 1);
                int idx111 = (k + 1) * dimX * dimY + (j + 1) * dimX + (i + 1);
                int idx011 = (k + 1) * dimX * dimY + (j + 1) * dimX + i;

                hex->GetPointIds()->SetId(0, idx000);
                hex->GetPointIds()->SetId(1, idx100);
                hex->GetPointIds()->SetId(2, idx110);
                hex->GetPointIds()->SetId(3, idx010);
                hex->GetPointIds()->SetId(4, idx001);
                hex->GetPointIds()->SetId(5, idx101);
                hex->GetPointIds()->SetId(6, idx111);
                hex->GetPointIds()->SetId(7, idx011);

                ugrid->InsertNextCell(hex->GetCellType(), hex->GetPointIds());
            }
        }
    }

    // 创建原始单元ID数组
    vtkNew<vtkIdTypeArray> originalCellIds;
    originalCellIds->SetNumberOfComponents(1);
    originalCellIds->SetName("vtkOriginalCellIds");
    originalCellIds->SetNumberOfTuples(numCells);
    
    // 使用并行方式设置原始单元ID
    vtkSMPTools::For(0, numCells,
        [&](vtkIdType begin, vtkIdType end)
        {
            for (vtkIdType cellId = begin; cellId < end; ++cellId)
            {
                originalCellIds->SetValue(cellId, cellId);
            }
        });
    
    ugrid->GetCellData()->AddArray(originalCellIds);
    return ugrid;
}

// 平面控件回调函数
void PlaneCallback(vtkObject* caller, long unsigned int eventId, void* clientData, void* callData)
{
    auto planeWidget = static_cast<vtkPlaneWidget*>(caller);
    auto plane = static_cast<vtkPlane*>(clientData);

    // 从平面控件获取当前的平面参数
    planeWidget->GetPlane(plane);

    // 更新提取几何的隐式函数
    extractGeometry->SetImplicitFunction(plane);
    extractGeometrySelected->SetImplicitFunction(plane);

    // 强制重新渲染
    renderer->GetRenderWindow()->Render();

    std::cout << "平面位置更新 - 当前裁剪后网格包含 "
              << extractGeometry->GetOutput()->GetNumberOfCells() << " 个体单元" << std::endl;
}

// 体单元拾取回调函数（现在已集成到CustomInteractorStyle中，这个函数可以移除）
/*
void PickerCallback(vtkObject* caller, long unsigned int vtkNotUsed(eventId),
                   void* vtkNotUsed(clientData), void* vtkNotUsed(callData))
{
    // 这个函数现在已经集成到CustomInteractorStyle::PerformCellPicking()中
}
*/

// 自定义交互样式类，用于处理键盘事件和鼠标事件
class CustomInteractorStyle : public vtkInteractorStyleTrackballCamera {
public:
    static CustomInteractorStyle* New();
    vtkTypeMacro(CustomInteractorStyle, vtkInteractorStyleTrackballCamera);

    // 设置拾取器
    void SetPicker(vtkHardwarePicker* picker) { this->Picker = picker; }

    virtual void OnMouseMove() override
    {
        // 检查是否在平面控件上点击
        if (planeWidget->GetEnabled()) {
            int* clickPos = this->GetInteractor()->GetEventPosition();

            // 尝试拾取平面控件
            planeWidget->SetInteractor(this->GetInteractor());
            // int picked = planeWidget->GetRepresentation()->ComputeInteractionState(clickPos[0], clickPos[1]);

            int picked = 0;

            // 如果没有点击在平面控件上，则进行体单元拾取
            if (picked == 0) // 0表示没有拾取到平面控件
            {
                this->PerformCellPicking();
                return; // 不调用父类方法，避免相机操作
            }
        } else {
            // 平面控件关闭时，直接进行体单元拾取
            this->PerformCellPicking();
            return;
        }

        // 调用父类方法处理相机操作
        vtkInteractorStyleTrackballCamera::OnLeftButtonDown();
    }

    void PerformCellPicking()
    {
        if (!this->Picker)
            return;

        int* clickPos = this->GetInteractor()->GetEventPosition();

        // 执行拾取
        this->Picker->Pick(clickPos[0], clickPos[1], 0, renderer);

        if (this->Picker->GetCellId() != -1) {
            std::cout << "拾取到的演员类: " << this->Picker->GetActor()->GetClassName() << std::endl;

            std::cout << "拾取到裁剪后polydata中的面单元ID: " << this->Picker->GetCellId() << std::endl;
            std::cout << "拾取位置: (" << this->Picker->GetPickPosition()[0] << ", "
                      << this->Picker->GetPickPosition()[1] << ", "
                      << this->Picker->GetPickPosition()[2] << ")" << std::endl;
                      
            // 获取裁剪后的网格
            vtkUnstructuredGrid* clippedGrid = extractGeometry->GetOutput();
            auto clippedcelldata = clippedGrid->GetCellData();
            if (clippedcelldata) {
                std::cout << "裁剪后的体网格包含 " << clippedcelldata->GetNumberOfArrays() << " 个数据数组" << std::endl;
                for (int i = 0; i < clippedcelldata->GetNumberOfArrays(); i++) {
                    auto array = clippedcelldata->GetArray(i);
                    if (array) {
                        std::cout << "数组 " << i << ": 名称 = " << array->GetName()
                                  << ", 元素数量 = " << array->GetNumberOfTuples() << std::endl;
                    }
                }
            } else {
                std::cout << "裁剪后的网格没有单元数据" << std::endl;
            }

            vtkDataSet* pickPoly = this->Picker->GetDataSet();
            if (pickPoly) {
                std::cout << "拾取到的多边形数据包含 " << pickPoly->GetNumberOfCells() << " 个单元" << std::endl;

                auto cell_data = pickPoly->GetCellData();
                if (cell_data) {
                    std::cout << "多边形数据包含 " << cell_data->GetNumberOfArrays() << " 个数据数组" << std::endl;
                    for (int i = 0; i < cell_data->GetNumberOfArrays(); i++) {
                        auto array = cell_data->GetArray(i);
                        if (array) {
                            std::cout << "数组 " << i << ": 名称 = " << array->GetName()
                                      << ", 元素数量 = " << array->GetNumberOfTuples() << std::endl;
                        }
                    }
                } else {
                    std::cout << "多边形数据没有单元数据" << std::endl;
                }
            } else {
                std::cout << "未拾取到多边形数据" << std::endl;
            }

            auto vtkOriginalCellIds = pickPoly->GetCellData()->GetArray("vtkOriginalCellIds");
            if (vtkOriginalCellIds) {
                vtkIdType originalId2 = vtkOriginalCellIds->GetTuple1(this->Picker->GetCellId());
                std::cout << "vtkOriginalCellIds 数组中的数据: " << originalId2 << std::endl;
            }
            else {
                std::cout << "未找到 vtkOriginalCellIds 数组" << std::endl;
            }

            vtkNew<vtkIdTypeArray> ids;
            ids->SetNumberOfComponents(1);
            ids->InsertNextValue(vtkOriginalCellIds->GetTuple1(this->Picker->GetCellId()));

            vtkNew<vtkSelectionNode> selectionNode;
            selectionNode->SetFieldType(vtkSelectionNode::CELL);
            selectionNode->SetContentType(vtkSelectionNode::INDICES);
            selectionNode->SetSelectionList(ids);

            vtkNew<vtkSelection> selection;
            selection->AddNode(selectionNode);

            vtkNew<vtkExtractSelection> extractSelection;
            extractSelection->SetInputData(0, originalGrid);
            extractSelection->SetInputData(1, selection);
            extractSelection->Update();

            // In selection
            vtkUnstructuredGrid* selected = vtkUnstructuredGrid::SafeDownCast(extractSelection->GetOutput());
            selectedGrid->ShallowCopy(selected); // 更新全局高亮网格

            std::cout << "\n";
            // 重新渲染
            renderer->GetRenderWindow()->Render();
        } else {
            std::cout << "未拾取到任何体单元" << std::endl;
        }
    }

    virtual void OnKeyPress() override
    {
        vtkRenderWindowInteractor* rwi = this->Interactor;
        std::string key = rwi->GetKeySym();

        if (key == "p" || key == "P") {
            // 切换平面控件的显示状态
            if (planeWidget->GetEnabled()) {
                planeWidget->Off();
                std::cout << "平面控件已隐藏 - 现在可以直接拾取体单元" << std::endl;
            } else {
                planeWidget->On();
                std::cout << "平面控件已显示 - 点击空白区域可拾取体单元" << std::endl;
            }
        } else if (key == "r" || key == "R") {
            selectedGrid->Initialize(); // 清空高亮显示
            renderer->GetRenderWindow()->Render();
            std::cout << "颜色已重置" << std::endl;
        }

        // 调用父类方法
        vtkInteractorStyleTrackballCamera::OnKeyPress();
    }

private:
    vtkHardwarePicker* Picker;
};

vtkStandardNewMacro(CustomInteractorStyle);

int main()
{
    system("chcp 65001");
    // 创建原始立方体网格
    originalGrid = CreateCubeGrid();
    selectedGrid = vtkSmartPointer<vtkUnstructuredGrid>::New();

    std::cout << "原始网格包含 " << originalGrid->GetNumberOfCells() << " 个体单元" << std::endl;

    // 创建平面对象
    plane = vtkSmartPointer<vtkPlane>::New();
    plane->SetOrigin(3.0, 3.0, 3.0); // 平面原点
    plane->SetNormal(1.0, 0.0, 0.0); // 平面法向量

    // 创建几何提取器（用于裁剪）
    extractGeometry = vtkSmartPointer<vtkExtractGeometry>::New();
    extractGeometrySelected = vtkSmartPointer<vtkExtractGeometry>::New();
    vtkNew<vtkUnstructuredGrid> originCopy;
	originCopy->ShallowCopy(originalGrid);
    extractGeometry->SetInputData(originCopy);
    extractGeometrySelected->SetInputData(selectedGrid);
    extractGeometry->SetImplicitFunction(plane);
    extractGeometrySelected->SetImplicitFunction(plane);
    extractGeometry->ExtractInsideOn(); // 提取平面法向量方向一侧的单元
    extractGeometrySelected->ExtractInsideOn();
    extractGeometry->ExtractBoundaryCellsOff(); // 不提取边界单元
    extractGeometrySelected->ExtractBoundaryCellsOff();

    // 初始化高亮演员与映射器
    selectedMapper = vtkSmartPointer<vtkDataSetMapper>::New();
    selectedMapper->SetInputConnection(extractGeometrySelected->GetOutputPort()); // 初始时没有高亮显示
    selectedMapper->SetRelativeCoincidentTopologyPolygonOffsetParameters(0, -1);
    selectedActor = vtkSmartPointer<vtkActor>::New();
    selectedActor->SetMapper(selectedMapper);
    selectedActor->GetProperty()->SetOpacity(1.0);
    selectedActor->GetProperty()->SetColor(1.0, 0.0, 0.0); // 红色高亮
    selectedActor->GetProperty()->EdgeVisibilityOn();
    selectedActor->GetProperty()->SetEdgeColor(0.0, 0.0, 0.0);

    auto to_polydata = vtkSmartPointer<vtkGeometryFilter>::New();
    to_polydata->SetInputConnection(extractGeometry->GetOutputPort());

    // 创建映射器
    mapper = vtkSmartPointer<vtkPolyDataMapper>::New();
    mapper->SetInputConnection(to_polydata->GetOutputPort());

    // 创建演员
    actor = vtkSmartPointer<vtkActor>::New();
    actor->SetMapper(mapper);
    actor->GetProperty()->SetOpacity(0.9);
    actor->GetProperty()->EdgeVisibilityOn(); // 显示边缘
    actor->GetProperty()->SetEdgeColor(0.0, 0.0, 0.0); // 黑色边缘

    // 创建渲染器
    renderer = vtkSmartPointer<vtkRenderer>::New();
    renderer->AddActor(actor);
    renderer->AddActor(selectedActor); // 添加高亮演员
    renderer->SetBackground(0.1, 0.1, 0.2); // 深蓝色背景

    // 创建渲染窗口
    vtkNew<vtkRenderWindow> renderWindow;
    renderWindow->AddRenderer(renderer);
    renderWindow->SetSize(1000, 800);
    renderWindow->SetWindowName("VTK体网格裁剪和拾取综合示例");

    // 创建交互器
    vtkNew<vtkRenderWindowInteractor> interactor;
    interactor->SetRenderWindow(renderWindow);

    // 创建体单元拾取器
    vtkNew<vtkHardwarePicker> picker;
    picker->PickFromListOn();
    picker->AddPickList(actor);
    //picker->SetTolerance(0.001);

    // 设置自定义交互样式
    vtkNew<CustomInteractorStyle> style;
    style->SetPicker(picker); // 将拾取器传递给交互样式
    interactor->SetInteractorStyle(style);

    // 创建平面控件
    planeWidget = vtkSmartPointer<vtkPlaneWidget>::New();
    planeWidget->SetInteractor(interactor);
    planeWidget->SetInputData(originalGrid);
    planeWidget->SetOrigin(3.0, 3.0, 3.0);
    planeWidget->SetNormal(1.0, 0.0, 0.0);
    planeWidget->SetResolution(10); // 平面控件的分辨率
    planeWidget->SetRepresentationToOutline(); // 显示为轮廓
    planeWidget->PlaceWidget();

    // 设置平面控件的外观
    planeWidget->GetPlaneProperty()->SetColor(1.0, 1.0, 0.0); // 黄色
    planeWidget->GetPlaneProperty()->SetOpacity(0.3); // 降低透明度，减少干扰
    planeWidget->GetHandleProperty()->SetColor(0.0, 1.0, 0.0); // 绿色手柄
    planeWidget->GetSelectedHandleProperty()->SetColor(1.0, 0.0, 1.0); // 选中时紫色

    // 创建平面控件回调函数
    vtkNew<vtkCallbackCommand> planeCallback;
    planeCallback->SetCallback(PlaneCallback);
    planeCallback->SetClientData(plane);

    // 连接回调函数到平面控件
    planeWidget->AddObserver(vtkCommand::InteractionEvent, planeCallback);

    // 启用平面控件
    planeWidget->On();

    // 设置相机位置以获得更好的视角
    renderer->GetActiveCamera()->SetPosition(15, 15, 15);
    renderer->GetActiveCamera()->SetFocalPoint(3, 3, 3);
    renderer->GetActiveCamera()->SetViewUp(0, 0, 1);
    renderer->ResetCameraClippingRange();

    // 开始渲染和交互
    renderWindow->Render();

    std::cout << "\n=== VTK体网格裁剪和拾取综合示例 ===" << std::endl;
    std::cout << "操作说明：" << std::endl;
    std::cout << "【平面裁剪操作】" << std::endl;
    std::cout << "- 拖拽黄色平面来改变裁剪位置" << std::endl;
    std::cout << "- 拖拽绿色手柄来调整平面方向" << std::endl;
    std::cout << "- 按 'P' 键显示/隐藏平面控件" << std::endl;
    std::cout << "\n【体单元拾取操作】" << std::endl;
    std::cout << "- 鼠标左键单击体单元可拾取并高亮显示" << std::endl;
    std::cout << "- 当平面控件显示时：点击空白区域（非平面控件）拾取体单元" << std::endl;
    std::cout << "- 当平面控件隐藏时：可以直接点击任意体单元拾取" << std::endl;
    std::cout << "- 按 'R' 键重置所有单元颜色" << std::endl;
    std::cout << "\n【视角控制】" << std::endl;
    std::cout << "- 鼠标右键拖拽：缩放" << std::endl;
    std::cout << "- 鼠标中键拖拽：平移" << std::endl;
    std::cout << "- 按 'q' 或关闭窗口退出程序" << std::endl;
    std::cout << "\n当前裁剪后网格包含 " << extractGeometry->GetOutput()->GetNumberOfCells() << " 个体单元" << std::endl;
    std::cout << "\n提示：平面控件已显示 - 点击空白区域可拾取体单元" << std::endl;

    interactor->Start();

    return 0;
}